
import sys
import time
import indexing.featureindex
import coverage.basecountreader

if __name__=="__main__":
    sample1 = sys.argv[1]
    sample2 = sys.argv[2]
    if len(sys.argv)<=3:
        path1 = "/archive10/bowtie/Full_Length_Results/Sample_Parameters/%s/Current" % sample1
        path2 = "/archive10/bowtie/Full_Length_Results/Sample_Parameters/%s/Current" % sample2
    else:
        path1 = sys.argv[3]
        path2 = sys.argv[4]
    coverageindex1 = "%s/%s.coverage_summary.INDEX" % (path1, sample1)
    coverageindex2 = "%s/%s.coverage_summary.INDEX" % (path2, sample2)
    snpresults1 = "%s/%s.snp-results.txt" % (path1, sample1)
    snpresults2 = "%s/%s.snp-results.txt" % (path2, sample2)
    ix1 = indexing.featureindex.FeatureIndex(coverageindex1)
    ix2 = indexing.featureindex.FeatureIndex(coverageindex2)
    reader1 = coverage.basecountreader.BaseCountReader(ix1)
    reader2 = coverage.basecountreader.BaseCountReader(ix2)

    candidates = {}

    MIN_COV = 4
    MIN_VAR = 2
    
    # Initial list contains those positions with >= 4 coverage and
    # >= 2 variant reads in either sample, as found in the SNP results.
    
    for line in open(snpresults1):
        cols = line.strip().split("\t")
        reads = int(cols[6])
        cov = int(cols[10])
        if reads>=MIN_VAR and cov>=MIN_COV:
            key = ("chr%s" % cols[2], long(cols[3]), cols[5])
            candidates[key] = True
    for line in open(snpresults2):
        cols = line.strip().split("\t")
        reads = int(cols[6])
        cov = int(cols[10])
        if reads>=MIN_VAR and cov>=MIN_COV:
            key = ("chr%s" % cols[2], long(cols[3]), cols[5])
            candidates[key] = True

    print >> sys.stderr, "# Initial read of %ld total candidates for %s vs %s ... downselecting." % (len(candidates), sample1, sample2)

    # Downselect rules out positions where the SNP results show
    # coverage < 4 in either sample (leaving positions where there
    # is insufficient coverage in a sample with no variant bases).

    for line in open(snpresults1):
        cols = line.strip().split("\t")
        reads = int(cols[6])
        cov = int(cols[10])
        if cov < MIN_COV:
            key = ("chr%s" % cols[2], long(cols[3]), cols[5])
            if candidates.has_key(key):
                del candidates[key]
    for line in open(snpresults2):
        cols = line.strip().split("\t")
        reads = int(cols[6])
        cov = int(cols[10])
        if cov < MIN_COV:
            key = ("chr%s" % cols[2], long(cols[3]), cols[5])
            if candidates.has_key(key):
                del candidates[key]

    candidateList = candidates.keys()
    candidateList.sort()
    
    print >> sys.stderr, "# Read %ld total candidates for %s vs %s." % (len(candidateList), sample1, sample2)
        
    n = 0
    t0 = time.time()
    for (chr, pos, base) in candidateList:
        n += 1
        if 0==(n%1000):
            dt = time.time()-t0
            rate = (1.0*n)/dt
            numRemaining = len(candidateList) - n
            timeRemaining = numRemaining / rate
            print >> sys.stderr, "# [%s vs %s] Checked %d of %d candidates in t=%s; remaining time ~ %f" % (sample1, sample2, n, len(candidateList), str(time.time()-t0), timeRemaining)
        f1 = reader1[(chr, pos)]
        f2 = reader2[(chr, pos)]
        genomic = "?"
        if f1:
            f1 = f1.getDetails()
            genomic = f1.genomic
        if f2:
            f2 = f2.getDetails()
            genomic = f2.genomic

        # Discard positions with insufficient coverage on either sample
        if (not f1) or f1.reads < MIN_COV or (not f2) or f2.reads < MIN_COV:
            continue
        
        if not f1.snp and not f2.snp:
            print >> sys.stderr, "# Warning... couldn't find SNP at expected position", (chr, pos)
            continue
        st1 = st2 = "0\t0\t0\t0"

        cts = {}
        if f1.snp:
            for xp in ["A", "C", "G", "T"]:
                # cts[xp]=getCountExcludingFirstAndLastThree(f1,xp)
                cts[xp] = f1.cts.get(xp,0)
        else:
            for xp in ["A", "C", "G", "T"]: cts[xp]=0
            cts[f1.genomic] = f1.reads
        st1 = "%d\t%d\t%d\t%d" % (cts["A"], cts["C"], cts["G"], cts["T"])
        var1 = cts[base]
        cov1 = cts["A"]+cts["C"]+cts["G"]+cts["T"]
        if cov1 < MIN_COV: continue
        percent1 = "%.2f" % (100.0*var1 / cov1)
        noncon1 = cov1 - cts[f1.genomic]
        nonconpct1 = "%.2f" % (100.0*noncon1 / cov1)
        
        cts = {}
        if f2.snp:
            for xp in ["A", "C", "G", "T"]:
                # cts[xp]=getCountExcludingFirstAndLastThree(f2,xp)
                cts[xp] = f2.cts.get(xp,0)
        else:
            for xp in ["A", "C", "G", "T"]: cts[xp]=0
            cts[f2.genomic] = f2.reads
        st2 = "%d\t%d\t%d\t%d" % (cts["A"], cts["C"], cts["G"], cts["T"])
        var2 = cts[base]
        cov2 = cts["A"]+cts["C"]+cts["G"]+cts["T"]
        if cov2 < MIN_COV: continue
        percent2 = "%.2f" % (100.0*var2 / cov2)
        noncon2 = cov2 - cts[f2.genomic]
        nonconpct2 = "%.2f" % (100.0*noncon2 / cov2)
        
        # Only print if either var1 or var2 is still >= 2
        if var1 < MIN_VAR and var2 < MIN_VAR: continue

        print "%s:%ld\t%s\t%s\t%s\t%s\t%d\t%d\t%d\t%d\t%d\t%d\t%s\t%s\t%s\t%s" % (chr, pos, genomic, base, percent1, percent2, var1, var2, cov1, cov2, noncon1, noncon2, nonconpct1, nonconpct2, st1, st2)
